# Untitled - By: luodongxu - 周一 6月 5 2023
from ulab import numpy as np
import random


def random_pick(some_list, probabilities):
    x = random.uniform(0,1)
    cumulative_probability = 0.0
    for item, item_probability in zip(some_list, probabilities):
         cumulative_probability += item_probability
         if x < cumulative_probability:
               break
    return item


def GetDistantMat(coordinates):
    """
    获取距离矩阵
    """
    n  = len(coordinates)
    res = np.zeros((n,n))
    for i in range(n):
        for j in range(i,n):
            lenth = np.sqrt((coordinates[i][0]-coordinates[j][0])**2 +(coordinates[i][1]-coordinates[j][1])**2)
            res[i,j]= res[j,i]=lenth
    return res;


def RandomPutAnts(AntNums:int,citysize:int):

    res = np.zeros((AntNums,citysize),dtype=np.int16)

    UnVisited:list[set] = []

    for i in range(AntNums):
        UnVisited.append(set(range(citysize)))
        res[i,0] = random.randrange(0,citysize,1)
        #res[i,0] = 0
        UnVisited[i].remove(res[i,0])
    return res,UnVisited



def ChosePath(Ettable:np.array,pheromonetable:np.array,nowcity:int,unvisited:set):
    n = Ettable.shape[0]

    tmp = Ettable[nowcity]*pheromonetable[nowcity]
    possible = []
    for citys in unvisited:
        if tmp[citys]<0:
            tmp[citys]=0.0
        possible.append(tmp[citys])

    allsum = np.sum(possible)
    if allsum == 0.0:
        Possiblelenth = len(possible)
        possible = [1.0/Possiblelenth]*Possiblelenth
    else:
        possible = [i/allsum for i in possible]

    return random_pick(list(unvisited),possible)

# TODO::还有bug
def AntsAlg(coordinates,AntNums:int,IterTurns:int,pheromoneDecree:float,Ke:float,Kp:float):
    n = len(coordinates)
    dtstanttable = GetDistantMat(coordinates)

    pheromonetable = np.ones((n,n))
    Ettable = 1/(dtstanttable**Ke + np.diag([float('inf')]*n))

    lengthbest = []  # 最佳路径长度
    lengthbestLenth = float('inf')

    for i in range(IterTurns):
        Ansplace,unvisted = RandomPutAnts(AntNums,n)
        LenthAnst = np.zeros(AntNums)
        for j in range(1,n):
            for Ants in range(AntNums):
                now = Ansplace[Ants][j-1]
                next = ChosePath(Ettable,pheromonetable,now,unvisted[Ants])
                unvisted[Ants].remove(next)
                Ansplace[Ants][j]=next
                pheromonetable[now][next] += Kp*pheromoneDecree/100
                LenthAnst[Ants]+= dtstanttable[now][next]
            pass
        for Ants in range(AntNums):
            now = 0
            next = Ansplace[Ants][n-1]
            pheromonetable[now][next] += Kp*pheromoneDecree/100
            LenthAnst[Ants]+= dtstanttable[now][next]
        if lengthbestLenth> np.min(LenthAnst):
            lengthbest = Ansplace[np.argmin(LenthAnst)]
            lengthbestLenth = np.min(LenthAnst)
        pheromonetable -=pheromoneDecree/100
        #print(pheromonetable)
    return lengthbest

"""
coord = [[0,0],[25,2],[26,5],[16,5],[5,10],[23,11],[20,13],[29,14],[0,0]]
res = AntsAlg(coord,30,20,1.2,1.2,1.1)
print(res)
"""
